"""
数据收集模块：从data/data_collector.py导入数据收集功能。
这是一个过渡模块，用于保持向后兼容性。
新代码应直接使用data/data_collector.py中的DataCollector类。
"""
import streamlit as st
import pandas as pd
import datetime
import logging
from pathlib import Path

# 从 app_setup 导入 logger
try:
    from .app_setup import logger
except ImportError:
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

# 导入数据库相关功能
try:
    from database.session import engine
    from database.models import Base
    DB_AVAILABLE = engine is not None
except ImportError as e:
    logger.error(f"无法导入数据库模块: {e}。数据加载将不可用。", exc_info=True)
    DB_AVAILABLE = False
    Base = None
    engine = None

# 导入新的DataCollector类
try:
    from database.data_collector import DataCollector
except ImportError as e:
    logger.error(f"无法导入DataCollector: {e}", exc_info=True)
    DataCollector = None

def create_db_tables():
    """检查并创建数据库表"""
    if engine:
        try:
            logger.info("检查并创建数据库表...")
            if DB_AVAILABLE and Base:
                Base.metadata.create_all(bind=engine)
                logger.info("数据库表检查/创建完成。")
            else:
                logger.warning("数据库不可用，无法创建表。")
        except Exception as e:
            logger.error(f"创建数据库表时出错: {e}", exc_info=True)
    else:
        logger.warning("数据库引擎未初始化，无法创建表。")

def load_initial_data():
    """加载初始数据到 session_state (使用数据库)"""
    if 'data_collector' not in st.session_state:
        if DataCollector:
            st.session_state.data_collector = DataCollector()
        else:
            logger.error("DataCollector类不可用，无法初始化数据收集器")
            st.session_state.data_collector = None
            return

    # 检查数据库是否可用，如果不可用且数据未加载，则显示警告并退出
    if not DB_AVAILABLE:
        if not st.session_state.get('data_loaded', False):
            st.warning("数据库未配置或连接失败，无法加载数据。请检查配置。")
            # 确保 session_state 中的数据结构存在，即使为空
            if 'indices_data' not in st.session_state:
                st.session_state.indices_data = {
                    "上证指数": pd.DataFrame(), "深证成指": pd.DataFrame(), "创业板指": pd.DataFrame(), "沪深300": pd.DataFrame(),
                    "stock_list": pd.DataFrame(columns=['code', 'name', 'industry'])
                }
            if 'stock_data' not in st.session_state:
                st.session_state.stock_data = {}
            st.session_state.data_loaded = False
        return

    # 数据库可用，继续加载逻辑
    if not st.session_state.get('data_loaded', False):
        try:
            data_collector = st.session_state.data_collector
            if data_collector is None:
                logger.error("数据收集器未初始化")
                return
                
            # 定义要加载的日期范围
            end_date = datetime.datetime.now().date()
            start_date = end_date - datetime.timedelta(days=730)  # 加载最近两年数据
            start_date_str = start_date.strftime("%Y%m%d")
            end_date_str = end_date.strftime("%Y%m%d")

            logger.info(f"开始从数据库加载初始数据 (日期范围: {start_date_str} 到 {end_date_str})...")

            # 获取指数数据
            sh_index = data_collector.get_index_data("000001", start_date_str, end_date_str)
            sz_index = data_collector.get_index_data("399001", start_date_str, end_date_str)
            cyb_index = data_collector.get_index_data("399006", start_date_str, end_date_str)
            hs300_index = data_collector.get_index_data("000300", start_date_str, end_date_str)

            # 股票列表
            stock_list = data_collector.get_stock_list()

            st.session_state.indices_data = {
                "上证指数": sh_index if not sh_index.empty else pd.DataFrame(),
                "深证成指": sz_index if not sz_index.empty else pd.DataFrame(),
                "创业板指": cyb_index if not cyb_index.empty else pd.DataFrame(),
                "沪深300": hs300_index if not hs300_index.empty else pd.DataFrame(),
                "stock_list": stock_list if not stock_list.empty else pd.DataFrame(columns=['code', 'name', 'industry'])
            }
            logger.info(f"指数数据加载完成。上证: {len(sh_index)} 行, 深证: {len(sz_index)} 行, 创业板: {len(cyb_index)} 行, 沪深300: {len(hs300_index)} 行, 列表: {len(stock_list)} 行")

            # 预加载常用个股数据
            common_stocks_codes = ["600519", "000858", "600036", "601318", "000333"]
            stock_data = {}
            for code in common_stocks_codes:
                try:
                    data = data_collector.get_daily_data(code, start_date_str, end_date_str)
                    if not data.empty:
                        stock_data[code] = data
                        logger.info(f"已预加载个股数据: {code} - {len(data)} 行")
                    else:
                        logger.warning(f"无法预加载个股数据: {code}")
                except Exception as e:
                    logger.error(f"加载{code}数据异常: {e}")
                    stock_data[code] = pd.DataFrame()

            st.session_state.stock_data = stock_data
            st.session_state.data_loaded = True
            logger.info("数据库初始数据加载完成。")

        except Exception as e:
            st.session_state.data_loaded = False
            logger.error(f"从数据库加载初始数据失败: {e}", exc_info=True)
            st.error(f"加载初始数据失败: {e}")
            # 初始化为空，避免后续代码出错
            st.session_state.indices_data = {
                "上证指数": pd.DataFrame(), "深证成指": pd.DataFrame(), "创业板指": pd.DataFrame(), "沪深300": pd.DataFrame(),
                "stock_list": pd.DataFrame(columns=['code', 'name', 'industry'])
            }
            st.session_state.stock_data = {}